/*	
*	File: udp-header.h
*	Date: 2011/11/07
*	User: Ronmaru
*	Desctiption: udp hader class
*/

//   0                         16                          31
//   +-------------------------+------------------------------+
//   |   source port UDP       |     destination port  UDP    |
//   +-------------------------+------------------------------+
//   |   message length  UDP   |     checksum UDP             |
//   +-------------------------+------------------------------+
//


#ifndef _UDP_HEADER_H_
#define _UDP_HEADER_H_

const int UDP_MAX_HEADER_LENGTH = 8;

class UDPheader
{
public:
  UDPheader() { std::fill(rep_, rep_ + sizeof(rep_), 0); }

  void SetSourcePort(unsigned short n) { Encode(0, 1, n); }
  void SetDestinationPort(unsigned short n) { Encode(2, 3, n); }
  void SetMessageLength(unsigned short n) { Encode(4, 5, n); }
  void SetChecksum(unsigned short n) { Encode(6, 7, n); }

  unsigned int GetSourcePort() const { return Decode(0, 1); }
  unsigned int GetDestinationPort() const { return Decode(2, 3); }
  unsigned int GetMessageLength() const { return Decode(4, 5); }
  unsigned int GetChecksum() const { return Decode(6, 7); }

  friend std::istream& operator>>(std::istream& is, UDPheader& header)
  { return is.read(reinterpret_cast<char*>(header.rep_), UDP_MAX_HEADER_LENGTH); }

  friend std::ostream& operator<<(std::ostream& os, const UDPheader& header)
  { return os.write(reinterpret_cast<const char*>(header.rep_), UDP_MAX_HEADER_LENGTH); }

private:
  unsigned short Decode(int a, int b) const
  { return (rep_[a] << 8) + rep_[b]; }

  void Encode(int a, int b, unsigned short n)
  {
    rep_[a] = static_cast<unsigned char>(n >> 8);
    rep_[b] = static_cast<unsigned char>(n & 0xFF);
  }

  unsigned char rep_[UDP_MAX_HEADER_LENGTH];
};

class UDPpseudoheader
{
public:
  UDPpseudoheader() { std::fill(rep_, rep_ + sizeof(rep_), 0); }

  std::vector<unsigned char> GetHeader() { return std::vector<unsigned char>(rep_, rep_ + sizeof(rep_)); }

  void SetSourceAddress(boost::asio::ip::address_v4 address) 
  {
    boost::asio::ip::address_v4::bytes_type bytes = address.to_bytes();
    rep_[0] = bytes[0];
    rep_[1] = bytes[1];
    rep_[2] = bytes[2];
    rep_[3] = bytes[3];
  }
  void SetDestinationAddress(boost::asio::ip::address_v4 address)
  {
    boost::asio::ip::address_v4::bytes_type bytes = address.to_bytes();
    rep_[4] = bytes[0];
    rep_[5] = bytes[1];
    rep_[6] = bytes[2];
    rep_[7] = bytes[3];
  }
  void SetMessageProtocol() { Encode(8, 9, 17); }
  void SetDatagramLength(unsigned short n) { Encode(10, 11, n); }

  boost::asio::ip::address_v4 GetSourceAddress() const
  {
    boost::asio::ip::address_v4::bytes_type bytes
      = { { rep_[0], rep_[1], rep_[2], rep_[3] } };
    return boost::asio::ip::address_v4(bytes);
  }
  boost::asio::ip::address_v4 GetDestinationAddress() const
  {
    boost::asio::ip::address_v4::bytes_type bytes
      = { { rep_[4], rep_[5], rep_[6], rep_[7] } };
    return boost::asio::ip::address_v4(bytes);
  }
  unsigned int GetMessageProtocol() { Decode(8, 9); }
  unsigned int GetDatagramLength() { Decode(10, 11); }

  friend std::istream& operator>>(std::istream& is, UDPpseudoheader& header)
  { return is.read(reinterpret_cast<char*>(header.rep_), 12); }

  friend std::ostream& operator<<(std::ostream& os, const UDPpseudoheader& header)
  { return os.write(reinterpret_cast<const char*>(header.rep_), 12); }

private:
  unsigned short Decode(int a, int b) const
  { return (rep_[a] << 8) + rep_[b]; }

  void Encode(int a, int b, unsigned short n)
  {
    rep_[a] = static_cast<unsigned char>(n >> 8);
    rep_[b] = static_cast<unsigned char>(n & 0xFF);
  }

  unsigned char rep_[12];
};

template <typename Iterator>
void ComputeChecksum(UDPheader& header, UDPpseudoheader &pseudo_header,
  Iterator body_begin, Iterator body_end)
{
  header.SetChecksum(0);
  unsigned long int sum = header.GetSourcePort() + header.GetDestinationPort() + header.GetMessageLength();

  std::vector<unsigned char> pseudo_header_str = pseudo_header.GetHeader();
  for (int i = 0; i < pseudo_header_str.size(); i += 2) {
    sum += static_cast<unsigned long int>(pseudo_header_str[i]) << 8;
    sum += static_cast<unsigned long int>(pseudo_header_str[i + 1]);
  }

  Iterator body_iter = body_begin;
  while (body_iter != body_end)
  {
    sum += (static_cast<unsigned char>(*body_iter++) << 8);
    if (body_iter != body_end)
      sum += static_cast<unsigned char>(*body_iter++);
  }

  sum = (sum >> 16) + (sum & 0xFFFF);
  sum += (sum >> 16);
  header.SetChecksum(static_cast<unsigned short>(~sum));
}


#endif //_UDP_HEADER_H_